#!/usr/bin/env python
'''

Created on 04/dic/2012

Project: UNIMORE IMAGE SEARCH
A simple image classification tool

@author:Arletti Luca, Bicocchi Nicola, Spina Andrea

'''

from config import Config

from nltk.corpus import wordnet as wn
import networkx as nx
import matplotlib.pyplot as plt

class Wordnet_Voter:
    
    def __init__(self):
        self.tree = nx.Graph()
        return
            
    def __add_vote(self, node):
        self.tree.add_node(node)
        try:
            self.tree.node[node]['votes'] += 1
        except KeyError:
            self.tree.node[node]['votes'] = 1
        return 
        
    def __connect_to(self, keyword, root='entity.n.01'):
        try:
            child = wn.synsets(keyword)[0]
            father = child.hypernyms()[0]
            self.__add_vote(child.name)
        except IndexError:
            self.problems.append(keyword)
            return
        
        distance = 0
        while True:
            try:
                father = child.hypernyms()[0]
            except IndexError:
                if child.name != root:
                    self.problems.append(child.name)
                    self.tree.add_edge(child.name, root)
                    self.__add_vote(root)
                return distance
            
            self.tree.add_edge(child.name, father.name)
            self.__add_vote(father.name)
                            
            distance += 1    
            child = father
            
        return distance
            
    def vote(self, keywords):
        self.tree.clear()
        self.keywords = keywords
        self.problems = []
        self.levels = {}
        self.best_node = None
        
        # build graph
        for keyword in keywords:
            self.__connect_to(keyword) 
            
        # select best node
        try:
            for node, distance in nx.shortest_path_length(self.tree, source='entity.n.01').items():
                self.tree.node[node]['distance'] = distance
        except KeyError:
            # missing entity.n.01 node. nothing to do
            return
            
        # organizes data in distance-based levels
        for node in self.tree.nodes():
            votes = self.tree.node[node]['votes']
            distance = self.tree.node[node]['distance']
            try:
                self.levels[distance].append({'name' : node, 'votes' : votes})
            except KeyError:
                self.levels[distance] = [{'name' : node, 'votes' : votes}]
        
        # searches for the best node
        for level in range(max(self.levels.keys()), -1, -1):
            for node in self.levels[level]:
                if node['votes'] >= (len(self.keywords) - len(self.problems)) * 0.5: 
                    break
            else:
                continue
            break
        
        self.best_node = node['name']
        return node['name']   

    def draw_graph(self, filename=None):
        import re
        import random
        
        if self.best_node is None:
            return 
        
        # prepare labels
        labels = {}
        for node_name, d in self.tree.nodes(data=True):
            labels[node_name] = '%s=%d' % (node_name, d['votes'])
        
        # prepare colors
        node_color = []
        limit = max(self.levels.keys())
        alpha = 1.2
        for n, d in self.tree.nodes(data=True):
            if n == self.best_node:
                # yellow
                node_color.append((1.0, 1.0, 0.0))
            else:
                # shades of red
                node_color.append((1.0, d['distance']/(limit * alpha), d['distance']/(limit * alpha)))
        
        # draw graph
        pos = nx.pygraphviz_layout(self.tree)
        nx.draw_networkx_nodes(self.tree, pos, node_color=node_color, node_size=12, linewidths=0.0, alpha=0.8)
        nx.draw_networkx_edges(self.tree, pos, width=0.1, alpha=0.5)
        nx.draw_networkx_labels(self.tree, pos, labels=labels, font_size=1.3, font_color='black')
        #plt.show()
        if filename is not None:
            plt.savefig(filename)
        plt.close()
        return
        
        

if __name__ == '__main__':
    wv = Wordnet_Voter()
    print wv.vote(['cat', 'dog', 'shark', 'butterfly'])
    wv.draw_graph('test.png')
    
    
        
        
        
        
        
        
        
        